package Spectrum

import (
	"encoding/json"
	"github.com/gorilla/websocket"
	"log"
	"math"
	"math/rand/v2"
	"net/http"
)

// WebSocket 消息类型
type SpectrumMsg struct {
	Type          string    `json:"type"` // "spectrum" 或 "demod"
	CenterFreq    float64   `json:"centerFreq,omitempty"`
	SampleRate    float64   `json:"sampleRate,omitempty"`
	Data          []float32 `json:"data,omitempty"`
	Constellation []float32 `json:"constellation,omitempty"` // [I,Q,I,Q...]
	Bits          string    `json:"bits,omitempty"`
	SNR           float32   `json:"snr,omitempty"`
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

type Hub struct {
	clients   map[*websocket.Conn]bool
	broadcast chan SpectrumMsg
}

func NewHub() *Hub {
	return &Hub{
		clients:   make(map[*websocket.Conn]bool),
		broadcast: make(chan SpectrumMsg, 100),
	}
}

func (h *Hub) Run() {
	for msg := range h.broadcast {
		data, _ := json.Marshal(msg)
		for c := range h.clients {
			if err := c.WriteMessage(websocket.TextMessage, data); err != nil {
				c.Close()
				delete(h.clients, c)
			}
		}
	}
}

func UpgradeWebSocket(h *Hub) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Println(err)
			return
		}
		h.clients[conn] = true
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				delete(h.clients, conn)
				conn.Close()
				break
			}
		}
	}
}

// ====== 生成模拟 FFT 测试数据（替代你的 IPP FFT） ======
func generateFFT() []float32 {
	N := 16384
	out := make([]float32, N)
	for i := 0; i < N; i++ {
		// 模拟两个窄带峰 + 噪声
		x1 := math.Exp(-math.Pow(float64(i-100)/6, 2)) * 30
		x2 := math.Exp(-math.Pow(float64(i-300)/8, 2)) * 25
		noise := rand.Float64() * 5

		amp := x1 + x2 + noise // 范围大约 0~30

		// 转换为 dB 范围 0 ~ -120
		db := -120 + amp*4
		if db > 0 {
			db = 0
		}
		out[i] = float32(db)
	}
	return out
}

type Broadcaster struct {
	input   chan []float32
	outputs []chan []float32
	done    chan struct{}
}

func NewBroadcaster(buffer int) *Broadcaster {
	return &Broadcaster{
		input:   make(chan []float32, buffer),
		outputs: make([]chan []float32, 0, 4),
		done:    make(chan struct{}),
	}
}

func (b *Broadcaster) Run() {
	for {
		select {
		case packet, ok := <-b.input:
			if !ok {
				b.closeAll()
				return
			}
			// 广播给所有活跃消费者
			for _, ch := range b.outputs {
				// 不阻塞！满就丢（保证生产者不卡）
				select {
				case ch <- packet:
				default:
				}
			}
		case <-b.done:
			return
		}
	}
}

func (b *Broadcaster) Subscribe() <-chan []float32 {
	ch := make(chan []float32, 8) // 每个消费者小缓冲
	b.outputs = append(b.outputs, ch)
	return ch
}

func (b *Broadcaster) Submit(data []float32) {
	select {
	case b.input <- data:
	case <-b.done:
	}
}

func (b *Broadcaster) Close() {
	close(b.done)
	close(b.input)
}

func (b *Broadcaster) closeAll() {
	for _, ch := range b.outputs {
		close(ch)
	}
}
